# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import json
import logging
import numpy

from coin.strategy.basis_strat.moving_average import MovingAverage, MovingObject
from coin.experimental.xguo.fruit.base.single_product import SingleProductPriceCalculator
from coin.experimental.xguo.fruit.util.fee import get_fee_rate
from coin.experimental.xguo.fruit.util.util import MovingWindowData


class SmmPricerConfig(object):
  def __init__(self):
    # numbers
    self.window_period = int(300 * 1e9)
    self.book_level = 5

    self.force_fire_threshold = 3 * 1e-4
    self.force_fire_cooldown = 5 * 1e9
    self.force_fire_cooldown_window = 5 * 1e9
    self.price_pull = 2 * 1e-4
    self.price_push = 2 * 1e-4
    self.force_fire_max_count = 3

    self.max_pos = None
    self.min_pos = None
    self.price_range = None

    # product
    self.trade_symbol = None
    self.ref_symbol = None
    self.get_pos = None

  def __str__(self):
    return json.dumps(self.__dict__, indent=2, default=str)


class SmmPricer(object):
  def __init__(self, config):
    self._logger = logging.getLogger("SmmPricer")
    self._price_calculators = {}
    self._timestamp = None
    self._config = config
    self._ma = MovingAverage(config.window_period)
    self._agg_buys = MovingObject(config.force_fire_cooldown_window)
    self._agg_sells = MovingObject(config.force_fire_cooldown_window)

    self._trade_ma = MovingWindowData(config.window_period)
    self._rate_ma = MovingWindowData(config.window_period)
    self._ref_ma = MovingWindowData(config.window_period)

    fee_bp = get_fee_rate(config.trade_symbol)[0]
    self._fee_rate = fee_bp * 1e-4

    products = (config.trade_symbol, config.ref_symbol)
    for prod in products:
      price_calculator = SingleProductPriceCalculator(prod)
      self._price_calculators[prod] = price_calculator

      if prod == config.trade_symbol:
        self._tpc = price_calculator
      elif prod == config.ref_symbol:
        self._rpc = price_calculator
      else:
        raise ValueError()
    self._logger.info("SmmPricer created for %s-%s", config.trade_symbol, config.ref_symbol)
    self._logger.info("config: %s", str(self._config))

  @property
  def trade_symbol(self):
    return self._config.trade_symbol

  @property
  def ref_symbol(self):
    return self._config.ref_symbol

  def get_pos(self):
    return self._config.get_pos()

  def get_norm_position(self):  # normalized to [-1, 1]
    pos_range = self._config.max_pos - self._config.min_pos

    if abs(pos_range) < 1e-5:
      return 0.0

    n_pos = 2 * (self.get_pos() - self._config.min_pos) / pos_range - 1
    return numpy.clip(n_pos, -1, 1)

  def is_ready(self):
    return self._ma.ready()

  def update_tick(self):
    self._tpc.renew_tick()
    self._rpc.renew_tick()

  def update(self, product, book):
    self._timestamp = book.timestamp
    self._price_calculators[product].update(book)

    if self._tpc.is_ready() and self._rpc.is_ready():
      # TODO(xguo) Handle no ask/bid/mid case.
      implied_rate = self._tpc.midp / self._rpc.midp
      self._ma.update(self._timestamp, implied_rate)

  def calc_edge(self, norm_pos, full_edge):
    pos = abs(norm_pos)
    # edge must cover fee_rate
    e2 = max(self._fee_rate, full_edge)
    e1 = (1 - pos) * e2 + self._fee_rate * pos

    assert e1 >= 0, e1
    assert e2 >= 0, e2

    # return buy_edge, sell_edge
    if norm_pos > 0:
      return e2, e1
    else:
      return e1, e2

  def gen_order(self):
    all_orders = {'product': str(self.trade_symbol)}
    rlist = [r[1] for r in self._ma._dq]
    p2, rate_median, p98 = numpy.percentile(
        rlist, (self._config.price_range[0], 50, self._config.price_range[1])) * 1e-6
    shift_edge1 = (p98 - rate_median) / rate_median
    shift_edge2 = (rate_median - p2) / rate_median
    shift_edge = max(shift_edge1, shift_edge2)

    ref = self._rpc
    trade = self._tpc

    norm_pos = self.get_norm_position()
    raw_mid_price = ref.midp * rate_median
    e1, e2 = self.calc_edge(norm_pos, shift_edge)
    raw_pass_buy_price = raw_mid_price * (1 - e1)
    raw_pass_sell_price = raw_mid_price * (1 + e2)

    if raw_pass_buy_price * (1 - self._config.force_fire_threshold) > trade.ask1:
      agg_buy_price = trade.get_prev_price(raw_pass_buy_price)
      agg_buy_qty = self._get_force_buy_qty(raw_pass_buy_price, agg_buy_price, trade)
      all_orders['agg_buy'] = dict(price=agg_buy_price, qty=agg_buy_qty)
      self._agg_buys.update(
          self._timestamp,
          (raw_pass_buy_price, agg_buy_price, (trade.ask1, trade.ask1_qty), agg_buy_qty))

    if raw_pass_sell_price * (1 + self._config.force_fire_threshold) < trade.bid1:
      agg_sell_price = trade.get_next_price(raw_pass_sell_price)
      agg_sell_qty = self._get_force_sell_qty(raw_pass_sell_price, agg_sell_price, self._tpc)
      all_orders['agg_sell'] = dict(price=agg_sell_price, qty=agg_sell_qty)
      self._agg_sells.update(
          self._timestamp,
          (raw_pass_sell_price, agg_sell_price, (trade.bid1, trade.bid1_qty), agg_sell_qty))

    pass_buy_price = trade.get_prev_price(min(raw_pass_buy_price, trade.ask1 - 1e-6))
    pass_sell_price = trade.get_next_price(max(raw_pass_sell_price, trade.bid1 + 1e-6))
    buy_price_pull = pass_buy_price * (1 + self._config.price_pull)
    sell_price_pull = pass_sell_price * (1 - self._config.price_pull)
    buy_price_push = pass_buy_price * (1 - self._config.price_push)
    sell_price_push = pass_sell_price * (1 + self._config.price_push)

    all_orders['mm'] = dict(
        pass_sell_price=pass_sell_price,
        pass_buy_price=pass_buy_price,
        sell_price_pull=sell_price_pull,
        buy_price_pull=buy_price_pull,
        sell_price_push=sell_price_push,
        buy_price_push=buy_price_push)
    all_orders['type'] = 'order'
    all_orders['raw_mid_price'] = raw_mid_price
    all_orders['p2'] = p2
    all_orders['p98'] = p98
    all_orders['rate_median'] = rate_median
    return all_orders

  def _get_agg_working_qty(self, mo):
    qty = 0.0
    for ts, obj in mo.iteritems():
      qty += obj[3]
    return qty

  def _get_force_buy_qty(self, raw_price, agg_price, pc):
    self._agg_buys.update(self._timestamp)
    size = self._agg_buys.size()
    qty = 0
    working_qty = 0
    avail_qty = 0

    if size >= self._config.force_fire_max_count:
      return 0

    if size <= 0:
      return 0

    last_ts = self._agg_buys.get_latest_ts()
    if self._timestamp - last_ts < self._config.force_fire_cooldown:
      return 0

    last_raw_price, last_agg_price, last_ba, _ = self._agg_buys.get_latest_data()
    # Order book not updated and agg price not improved.
    if (numpy.isclose(last_ba[0], pc.ask0) and
        numpy.isclose(last_ba[1], pc.ask0_qty) and agg_price <= last_agg_price + 1e-6):
      return 0

    # Agg raw price moving lower, do not fire.
    if raw_price < last_raw_price:
      return 0

    working_qty = self._get_agg_working_qty(self._agg_buys)
    for p, q in pc.get_ask_array(self._config.book_level):
      if p > raw_price:
        break
      avail_qty += q
    vacancy_qty = self._config.max_pos - self.get_pos() - working_qty
    qty = min(vacancy_qty, self._config.lot_size, avail_qty)
    return qty

  def _get_force_sell_qty(self, raw_price, agg_price, pc):
    self._agg_sells.update(self._timestamp)
    size = self._agg_sells.size()
    qty = 0
    working_qty = 0
    avail_qty = 0

    if size >= self._config.force_fire_max_count:
      return 0

    if size <= 0:
      return 0

    last_ts = self._agg_sells.get_latest_ts()
    if self._timestamp - last_ts < self._config.force_fire_cooldown:
      return 0

    last_raw_price, last_agg_price, last_bb, _ = self._agg_sells.get_latest_data()
    # Order book not updated and agg price not improved.
    if (numpy.isclose(last_bb[0], pc.bid0) and
        numpy.isclose(last_bb[1], pc.bid0_qty) and agg_price >= last_agg_price - 1e-6):
      return 0

    # Agg raw price moving higher, do not fire.
    if raw_price > last_raw_price:
      return 0

    working_qty = self._get_agg_working_qty(self._agg_sells)
    for b, q in pc.get_bid_array(self._config.book_level):
      if b < raw_price:
        break
      avail_qty += q
    vacancy_qty = self.get_pos() - self._config.min_pos - working_qty
    qty = min(vacancy_qty, self._config.lot_size, avail_qty)
    return qty
